#include "fkhook.h"

#include <memoryapi.h>
#include <winnt.h>
#include <winsock2.h>
#include <vector>

#include "tinyhook.h"

using sendto_func = int WINAPI(SOCKET, const char *, int, int, const sockaddr *, int);
using select_func = int WINAPI(int, fd_set *, fd_set *, fd_set *, const TIMEVAL *);
using recvfrom_func = int WINAPI(SOCKET, char *, int, int, sockaddr *, int *);

static TH_Info sendto_hook;
static TH_Info select_hook;
static TH_Info recvfrom_hook;
static sendto_func* _sendto = NULL;
static select_func* _select = NULL;
static recvfrom_func* _recvfrom = NULL;
static SOCKET origin_sock = 0, fake_sock = 0;
static std::vector<SOCKET> socks;

// 枚举当前所有可用网卡的IPv4地址
static const std::vector<in_addr> &enum_addr()
{
    static std::vector<in_addr> list;
    hostent *phost = gethostbyname(""); // 获取本机网卡
    if (phost) {
        char **ppc = phost->h_addr_list; // 获取地址列表
        int addr_num = 0;
        while (*ppc)
            addr_num++, ppc++; // 获取地址个数
        if (addr_num == list.size()) // 数量相同直接返回
            return list;
        ppc = phost->h_addr_list;
        list.clear();
        // 遍历列表添加到容器
        while (*ppc) {
            in_addr addr;
            memcpy(&addr, *ppc, sizeof(in_addr));
            list.push_back(addr);
            ppc++;
        }
    }
    return list;
}

// hook后替换的函数
static int WINAPI fake_sendto(SOCKET s, const char *buf, int len, int flags, const sockaddr *to, int tolen)
{
    sockaddr_in *toaddr = (sockaddr_in *)to;
    if (toaddr->sin_addr.S_un.S_addr != INADDR_BROADCAST)
        return _sendto(s, buf, len, flags, to, tolen); // 非广播直接原样发送
    int result = -1;
    origin_sock = s; // 暂存这个socket
    const std::vector<in_addr> &list = enum_addr();
    if (socks.size() != list.size()) {
        sockaddr_in addr_self;
        addr_self.sin_family = AF_INET;
        int namelen = sizeof(sockaddr_in);
        getsockname(s, (sockaddr *)&addr_self, &namelen); // 获取原sockaddr
        if (addr_self.sin_port == 0) {
            // 如果没有端口号，先原样发送，这样系统才会分配一个端口号
            result = _sendto(s, buf, len, flags, to, tolen);
            getsockname(s, (sockaddr *)&addr_self, &namelen); // 重新获取
        }
        for (int i = 0; i < socks.size(); i++)
            closesocket(socks[i]);
        socks.clear();
        for (int i = 0; i < list.size(); i++) {
            addr_self.sin_addr = list[i]; // 把新的地址换上去，然后绑定
            SOCKET sock = socket(AF_INET, SOCK_DGRAM, 0);
            BOOL opt = TRUE;
            setsockopt(sock, SOL_SOCKET, SO_BROADCAST, (char *)&opt, sizeof(BOOL)); // 广播
            setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, (char *)&opt, sizeof(BOOL)); // 重用地址端口
            if (bind(sock, (sockaddr *)&addr_self, sizeof(sockaddr)) == SOCKET_ERROR) {
                closesocket(sock);
                continue;
            } // 绑定到地址端口
            socks.push_back(sock);
        }
    }
    // 向列表中的每一个地址转发广播
    for (int i = 0; i < socks.size(); i++)
        result = _sendto(socks[i], buf, len, flags, to, tolen);
    return result;
}

static int WINAPI fake_select(int n, fd_set *rd, fd_set *wr, fd_set *ex, const TIMEVAL *timeout)
{
    if (rd && rd->fd_count == 1 && origin_sock == rd->fd_array[0]) {
        fd_set fds;
        FD_ZERO(&fds);
        for (int i = 0; i < socks.size(); i++)
            FD_SET(socks[i], &fds);
        int r = _select(0, &fds, NULL, NULL, timeout);
        if (r > 0) {
            fake_sock = fds.fd_array[0];
            return fds.fd_count;
        }
        fake_sock = 0;
        rd->fd_count = 0;
        return 0;
    }
    return _select(n, rd, wr, ex, timeout);
}

static int WINAPI fake_recvfrom(SOCKET s, char *buf, int len, int flags, sockaddr *from, int *fromlen)
{
    if (s == origin_sock && fake_sock != 0)
        return _recvfrom(fake_sock, buf, len, flags, from, fromlen);
    return _recvfrom(s, buf, len, flags, from, fromlen);
}

void hook()
{
    HMODULE h = GetModuleHandle(TEXT("ws2_32.dll"));
    void *p_sendto = (void*)GetProcAddress(h, "sendto");
    void *p_select = (void*)GetProcAddress(h, "select");
    void *p_recvfrom = (void*)GetProcAddress(h, "recvfrom");
    BYTE *padding = (BYTE*)TH_GetModulePadding(h);
    DWORD old_protect;
    constexpr int padding_size = 16 * 3;
    VirtualProtect(padding, padding_size, PAGE_EXECUTE_READWRITE, &old_protect);
    TH_Init(&sendto_hook, p_sendto, (void*)&fake_sendto, padding);
    TH_Init(&select_hook, p_select, (void*)&fake_select, padding + 16);
    TH_Init(&recvfrom_hook, p_recvfrom, (void*)&fake_recvfrom, padding + 32);
    VirtualProtect(padding, padding_size, PAGE_EXECUTE_READ, &old_protect);
    TH_GetDetour(&sendto_hook, (void**)&_sendto);
    TH_GetDetour(&select_hook, (void**)&_select);
    TH_GetDetour(&recvfrom_hook, (void**)&_recvfrom);
    TH_Hook(&sendto_hook);
    TH_Hook(&select_hook);
    TH_Hook(&recvfrom_hook);
}

void unhook()
{
    TH_Unhook(&sendto_hook);
    TH_Unhook(&select_hook);
    TH_Unhook(&recvfrom_hook);
    for (int i = 0; i < socks.size(); i++)
        closesocket(socks[i]);
    socks.clear();
}
